#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-09-18

"""
【节点名称】：
    MOTLoaderNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    无
【发布类型】：
    sensor_msgs::CompressedImage （输出图像）
【构造参数说明】：
    parameter_file (str): 全局参数文件（节点参数直接从主节点传过来，所以可以不用传参数文件了）
【节点参数】：
    mot_folder_path (str): MOT测试文件夹地址
【备注】：
    无
"""
import os
import threading
import cv2
import uuid
from typing import Union
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, QoS
from spirecv.base.BaseNode import BaseNode


class MOTLoaderNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown_emit: bool = True,
        mot_folder_path = ''
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=True
        )
        self.sms_shutdown_emit = sms_shutdown_emit
        self.client_id = str(uuid.uuid4()).replace('-', '_')
        self.img_i_queue = Queue()
        self._next_reader = Subscriber(
            '/' + job_name + '/launch_next', 'std_msgs::Boolean', self.launch_next,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._image_writer = Publisher(
            '/' + job_name + '/sensor/image_raw', 'sensor_msgs::CompressedImage',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self.image_fns = []
        self.img_len = []
        # self.mot_folder_path = self.get_param("mot_folder_path", os.path.expanduser('~')+"/mot20/train/")
        self.mot_folder_path = mot_folder_path
        self.img_len_num = 0
        self.frame_id = 0
        
        mot_files_and_directories = sorted(os.listdir(self.mot_folder_path), key=lambda x: (len(x), x), reverse=False)
        for seq_name in mot_files_and_directories:
            # print(seq_name)
            self.image_dir = self.mot_folder_path + seq_name + "/img1"
            self.img_len.append(len(os.listdir(self.image_dir)))
            files_and_directories = os.listdir(self.image_dir)
            # 图片整序
            img_paths = sorted(files_and_directories, key=lambda x: (len(x), x), reverse=False)
            # print(img_paths)
            for img_name in img_paths:
                if os.path.splitext(img_name)[-1] in ['.jpg', '.jpeg', '.png']:
                    self.image_fns.append(os.path.join(self.image_dir, img_name))
                    self.img_i = 0
        self.img_num = self.img_len[0]
        # print(len(self.image_fns))
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_writer.kill()
        self._next_reader.kill()

    def launch_next(self, msg: dict = None):
        if (isinstance(msg, dict) and msg['data']) or msg is None:
            if self.img_i < len(self.image_fns):
                #time.sleep(0.0417)
                self.img_i_queue.put(self.img_i)
                self.img_i += 1


    def run(self):
        while self.is_running():
            img_i = self.img_i_queue.get(block=True)
            if img_i is None:
                break
            if self.frame_id==self.img_num:#当发完一个文件，就重新imgid计数
                self.frame_id = 1
                self.img_len_num+=1
                self.img_num = self.img_len[self.img_len_num]
                #print(self.img_num)
            else:
                self.frame_id+=1

            img = cv2.imread(self.image_fns[img_i])
            
            msg = cvimg2sms(img)
            msg['img_id'] = self.frame_id
            msg['img_total'] = len(self.image_fns)
            msg['file_name'] = os.path.basename(self.image_fns[img_i])
            msg['client_id'] = self.client_id
            self._image_writer.publish(msg)
        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    node = MOTLoaderNode('detection')
